#!/usr/bin/env python
# -*- coding: utf-8 -*-

# usbmanager.py
#
# Copyright (c) 2008 Magnun Leno da Silva
#
# Author: Magnun Leno da Silva <magnun.leno@gmail.com>
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public License
# as published by the Free Software Foundation; either version 2 of
# the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
# USAA.

__version__ = "0.9"
__author__ = "Magnun Leno da Silva <magnun.leno@gmail.com>"

import gobject
import gtk
import dbus
import subprocess
import pygtk
pygtk.require('2.0')
from os import popen3, getuid
from sys import argv
from dbus.mainloop.glib import DBusGMainLoop


def ShowMessage(parent, text, type=0):
    '''
        A simple message box used to create some output. 
    '''
    if type is 0:
        type = gtk.MESSAGE_ERROR
    elif type is 1:
        type = gtk.MESSAGE_WARNING
    elif type is 2:
        type = gtk.MESSAGE_QUESTION
    else:
        type = gtk.MESSAGE_INFO
        
    dialog = gtk.MessageDialog(parent, gtk.DIALOG_DESTROY_WITH_PARENT, type,
                               gtk.BUTTONS_OK, text)
    dialog.run()
    dialog.destroy()
    
class Device:
    def __init__(self, volume, storage, usbdevices):
        self.volume = volume
        self.storage = storage
        self.devices = usbdevices
        self.bus = self.devices.bus
        
        self.label = None
        self.vendor = None
        self.size = None
        self.is_mounted = None
        self.read_only = None
        self.mount_point = None
        self.block_device = None
        self.fs_type = None
        self.fs_version = None
        
        self.udi = volume.GetProperty("info.udi")
        obj = self.bus.get_object ("org.freedesktop.Hal", self.udi)
        self.get_info()
        obj.connect_to_signal("PropertyModified", self.property_modified)
        
    def get_info(self):
        # Label
        self.label = str(self.volume.GetProperty("volume.label"))
        
        # Basic info
        self.vendor = str(self.storage.GetProperty("storage.vendor"))
        self.size = self.convert_bytes(self.volume.GetProperty("volume.size"))
        self.is_mounted = int(self.volume.GetProperty("volume.is_mounted"))
        self.read_only = int(self.volume.GetProperty("volume.is_mounted_read_only"))
        
        #Advanced info
        self.mount_point = str(self.volume.GetProperty("volume.mount_point"))
        self.block_device = str(self.volume.GetProperty("block.device"))
        self.fs_type = str(self.volume.GetProperty ('volume.fstype'))
        self.fs_version = str(self.volume.GetProperty ('volume.fsversion'))
    
    def property_modified(self, key, was_added=None, was_removed=None):
        self.devices.refresh_properties()
        
    def convert_bytes(self, size, base=0):
        if type(size) is int:
            size = float(size)
            
        if size >= 1024.0:
            base = base + 1
            size = size/1024.0
            return self.convert_bytes(size, base)
        else:
            return "%.2lf"%round(size,2) + " " + ['Bytes','KBytes','MBytes','GBytes','PBytes','EBytes','ZBytes','YBytes','BBytes'][base]
        
    def mount(self):
        opt = ['uid='+str(getuid())]
        obj = self.bus.get_object("org.freedesktop.Hal", self.udi)
        try:
            obj.Mount(self.label, '', opt, dbus_interface="org.freedesktop.Hal.Device.Volume")
        except dbus.DBusException, msg:
            ShowMessage(None, 'Erro while mounting '+self.block_device, 0)
            return 0
        return 1
    
    def umount(self):
        obj = self.bus.get_object("org.freedesktop.Hal", self.udi)
        try:
            obj.Unmount([''], dbus_interface="org.freedesktop.Hal.Device.Volume")
        except dbus.DBusException, msg:
            ShowMessage(None, 'Erro while umounting '+self.block_device, 0)
            return 0
        return 1

    def set_label(self, label):
        # It can be don with:
        # self.volume.SetPropertyString('volume.label', label)
        # But need to be running as root
        set_property = None
        ret = popen3("whereis hal-set-property")
        
        try:
            set_property = ret[1].readlines()[0].split(' ')[1]
        except:
            ShowMessage(None, "Unable to find the command 'hal-set-property'\nThis command is necessary to change the volume label")
            return 0
            
        comm = set_property+' --udi '+self.udi+' --key volume.label --string '+label
        msg = '-m "Changing label from '+self.block_device+'"'
        
        try:
            proc = subprocess.Popen(shell = False,
                                    stdin = subprocess.PIPE,
                                    stdout = subprocess.PIPE,
                                    stderr = subprocess.PIPE,
                                    args = ["gksudo", msg, comm])
            (output, error) = proc.communicate()
            if len(error) is not 0:
                ShowMessage(None, "Could not find the hal-set-property command")
                return 0
            proc.wait()
        except:
            ShowMessage(None, "Could not find the hal-set-property command")
            return 0
        return 1

class USBDevices:
    def __init__(self):
        self.devices = {}
        # Starting DBUS
        self.bus = dbus.SystemBus()
        self.hal_manager_obj = self.bus.get_object("org.freedesktop.Hal", "/org/freedesktop/Hal/Manager")
        self.hal_manager = dbus.Interface(self.hal_manager_obj, "org.freedesktop.Hal.Manager")
        
        
        # Windows 'containers'
        self.main = None
        self.properties = None
       
    def run(self, block_device=None):
       # Filter block_device content 
        if type(block_device) is not str and block_device is not None:
            ShowMessage(None, "Please specify a valid block device")
            raise TypeError, "block_device must be a valid string"
            exit()
        if block_device is not None and len(block_device) is 0:
            block_device = None
        
        if block_device is None:
            # main window
            self.hal_manager.connect_to_signal("DeviceAdded", self.new_device)
            self.hal_manager.connect_to_signal("DeviceRemoved", self.refresh)
            self.list_all_devices()
            self.main = MainWindow(self)
            self.main.run()
        else:
            # Only properties
            self.list_device(block_device)
            self.properties = PropertiesWindow(self.device[0][1], self)
            self.properties.run()
        
        exit(1)
        
        
    def list_all_devices(self):
        self.devices = {}
        self.volume_udi_list = self.hal_manager.FindDeviceByCapability('volume')
        for volume_udi in self.volume_udi_list:
            # DBUS Volume
            volume_obj = self.bus.get_object ("org.freedesktop.Hal", volume_udi)
            volume = dbus.Interface(volume_obj, "org.freedesktop.Hal.Device")
            
            # DBUS storage
            storage_udi = volume.GetProperty ('block.storage_device')
            storage_obj = self.bus.get_object ("org.freedesktop.Hal", storage_udi)
            storage = dbus.Interface(storage_obj, "org.freedesktop.Hal.Device")
            
            # Only USB
            if storage.GetProperty("storage.bus") == 'usb':
                self.devices[str(volume.GetProperty("volume.label"))] = Device(volume, storage, self)
    
    def list_device(self, block_device, show=True):
        self.devices = {}
        self.volume_udi_list = self.hal_manager.FindDeviceByCapability('volume')
        for volume_udi in self.volume_udi_list:
            # DBUS Volume
            volume_obj = self.bus.get_object ("org.freedesktop.Hal", volume_udi)
            volume = dbus.Interface(volume_obj, "org.freedesktop.Hal.Device")
            
            if str(volume.GetProperty("block.device")) == block_device:
                # DBUS storage
                storage_udi = volume.GetProperty ('block.storage_device')
                storage_obj = self.bus.get_object ("org.freedesktop.Hal", storage_udi)
                storage = dbus.Interface(storage_obj, "org.freedesktop.Hal.Device")
                
                # Only USB
                if storage.GetProperty("storage.bus") == 'usb':
                    self.devices[str(volume.GetProperty("volume.label"))] = Device(volume, storage, self)
        
        if len(self.devices) is 0:
            if show:
                ShowMessage(None, "Device %s not found or it's not a manageable device"%(block_device))
            return False
        return True
    
    def get_device(self, name):
        for device_name, device in self:
            if device_name == name:
                return device
        return None
    
    def refresh(self, volume_udi=None):
        self.list_all_devices()
        try:
            self.main.refresh()
            if self.properties is not None:
                showing_device = self.properties.device.label
                if self.get_device(showing_device) is None:
                    self.properties.window.destroy()
                    self.properties = None
        except:
            pass

    def refresh_properties(self):
        if self.properties is not None:
            self.properties.refresh()
        if self.main is not None:
            self.main.refresh_properties()
    
    def new_device(self, volume_udi):
        volume_obj = self.bus.get_object ("org.freedesktop.Hal", volume_udi)
        volume = dbus.Interface(volume_obj, "org.freedesktop.Hal.Device")
        
        if not volume.QueryCapability("volume"):
            return
            
        storage_udi = volume.GetProperty ('block.storage_device')
        storage_obj = self.bus.get_object ("org.freedesktop.Hal", storage_udi)
        storage = dbus.Interface(storage_obj, "org.freedesktop.Hal.Device")
        
        if storage.GetProperty("storage.bus") != 'usb':
            return
            
        self.refresh()
        
    def __len__(self):
        return len(self.devices)
        
    def __getitem__(self, key):
        return self.devices.items()[key]


class PropertiesWindow:
    def __init__(self, device, usbdevices, slave_window=False):
        self.device = device
        self.devices = usbdevices
        self.auto_refresh = False
        # Main Window
        self.window = gtk.Window(gtk.WINDOW_TOPLEVEL)
        if slave_window is True:
            self.window.set_modal(True)
            
        self.window.connect("destroy", self.destroy)
        self.window.set_border_width(10)
        self.window.set_size_request(300,240)
        self.window.set_title('USB Storage Device Properties')
        self.window.show()
        
        # Main VBox
        self.vbox = gtk.VBox()
        self.window.add(self.vbox)
        
        ##### Device Name #####
        self.namebox = gtk.HBox()
        self.nameLabel = gtk.Label('Volume Label:')
        self.nameLabel.set_alignment(0.1,0.5)
        self.nameEntry = gtk.Entry()
        self.nameEntry.set_text(device.label)
        self.namebox.pack_start(self.nameLabel)
        self.namebox.pack_start(self.nameEntry)
        self.vbox.pack_start(self.namebox)
        
        ##### Details #####
        # Detail expander
        self.detExp = gtk.Expander('<b>Details</b>')
        self.detExp.set_use_markup(True)
        self.detExp.set_expanded(True)
        
        # Detail Table
        self.detTable = gtk.Table(4,2)
        detLabels = [('Device Vendor :', device.vendor),
                     ('Volume Size :', device.size),
                     ('Is Mounted :', device.is_mounted),
                     ('Read Only :', device.read_only)]
        self.details = []
        max_len = max([len(det[0]) for det in detLabels]) + 5
        # Add details labels
        for i, det in enumerate(detLabels):
            # Label
            label = gtk.Label(det[0].rjust(max_len))
            label.set_alignment(0.1,0.5)
            
            # Entry or Checkbox
            if type(det[1]) is int:
                entry = gtk.CheckButton()
                entry.set_active(det[1])
            else:
                entry = gtk.Entry()
                entry.set_text(det[1].strip())
            if i is 2:
                entry.connect('toggled', self.toggle)
            else:
                entry.set_sensitive(0)
            self.details.append((label, entry))
            
            self.detTable.attach(label, 0, 1, i, i+1)
            self.detTable.attach(entry, 1, 2, i, i+1)
        
        self.detExp.add(self.detTable)
        self.vbox.pack_start(self.detExp, padding=5)
        
        ##### Advanced #####
        # Advanced expander
        self.advExp = gtk.Expander('<b>Advanced</b>')
        self.advExp.set_use_markup(True)
        self.advExp.set_expanded(False)
        self.advExp.connect("activate", self.expanded)
        
        # Avanced Table
        self.advTable = gtk.Table(4,2)
        self.advLabels = [('Block Device :', device.block_device),
                          ('Mount Point :', device.mount_point),
                          ('File System :', device.fs_type),
                          ('FS Version :', device.fs_version)]
        self.advanced = []
        max_len = max([len(det[0]) for det in self.advLabels]) + 5
        
        # Add advanced labels
        for i, adv in enumerate(self.advLabels):
            # Label
            label = gtk.Label(adv[0].rjust(max_len))
            label.set_alignment(0.1,0.5)
            self.advTable.attach(label, 0, 1, i, i+1)
            # Entry
            entry = gtk.Entry()
            entry.set_text(adv[1].strip())
            entry.set_sensitive(0)
            self.advTable.attach(entry, 1, 2, i, i+1)
            self.advanced.append((label, entry))
            
        self.advExp.add(self.advTable)
        self.vbox.pack_start(self.advExp, padding=5)
        
        ##### Buttons #####
        self.btnBox = gtk.HButtonBox()
        self.btnBox.set_layout(gtk.BUTTONBOX_SPREAD)
        self.okBtn = gtk.Button("OK", "gtk-ok")
        self.okBtn.connect("clicked", self.apply)
        self.cancelBtn = gtk.Button("Cancel", "gtk-cancel")
        self.cancelBtn.connect("clicked", self.destroy)
        self.btnBox.pack_start(self.okBtn)
        self.btnBox.pack_start(self.cancelBtn)
        self.vbox.pack_start(self.btnBox, padding=5)
        
        self.window.show_all()
    
    def expanded(self, exp):
        if not exp.get_expanded():
            # Is being closed, need to enlarge
            actual_size = self.window.get_size() 
            if actual_size[1] < 332:
                self.window.set_size_request(300, 332)
                self.window.resize(300, 332)
        else:
            # Is being expanded, need to reduce
            self.window.set_size_request(300, 240)
            self.window.resize(300, 240)
    
    def apply(self, widget):
        if self.device.label != self.nameEntry.get_text():
            self.device.set_label(self.nameEntry.get_text())
        self.window.destroy()
        
    def toggle(self, widget):
        if self.auto_refresh:
            return
        if widget.get_active(): # Mount
            self.device.mount()
        else: # Umount
            self.device.umount()
        #self.refresh()
        #if self.mainwindow:
        #    self.mainwindow.refresh()
    
    def refresh(self):
        self.auto_refresh = True
        self.device.get_info()
        self.nameEntry.set_text(self.device.label)
        self.details[0][1].set_text(self.device.vendor.strip())
        self.details[1][1].set_text(self.device.size.strip())
        self.details[2][1].set_active(self.device.is_mounted)
        self.details[3][1].set_active(self.device.read_only)
        
        self.advanced[0][1].set_text(self.device.block_device)
        self.advanced[1][1].set_text(self.device.mount_point)
        self.advanced[2][1].set_text(self.device.fs_type)
        self.advanced[3][1].set_text(self.device.fs_version)
        self.auto_refresh = False
    
    def destroy(self, widget, data=None):
        self.window.hide()
        self.devices.properties = None
        self.window.destroy()
        gtk.main_quit()

    def run(self):
        gtk.main()
    

class MainWindow:
    def __init__(self, usbdevices):
        self.devices = usbdevices
        
        self.window = gtk.Window(gtk.WINDOW_TOPLEVEL)
        self.window.connect("destroy", self.destroy)
        self.window.set_border_width(10)
        self.window.set_size_request(300,240)
        self.window.set_title('USB Storage Devices')
        self.window.show()
        self.devices.main = self.window
        
        # Main VBox
        self.vbox = gtk.VBox()
        self.window.add(self.vbox)
        
        # TreeView
        self.devicesView = gtk.TreeView()
        self.add_text_column("Label", 0)
        self.add_text_column("Block Device", 1)
        self.add_toggle_column("Mounted", 2)
        
        
        # ListStore
        self.devicesList = gtk.ListStore(str, str, gobject.TYPE_BOOLEAN)
        self.devicesView.set_model(self.devicesList)
        self.vbox.pack_start(self.devicesView)
        
        if len(self.devices) is not 0:
            for name, device in self.devices:
                #device.get_info()
                self.devicesList.append([name, device.block_device, device.is_mounted])
        
        # Signals
        self.devicesView.connect("row-activated", self.row_click)
        
        self.window.show_all()
        
    def add_text_column(self, title, id):
        column = gtk.TreeViewColumn(title, gtk.CellRendererText(), text=id)
        column.set_resizable(True)
        column.set_sort_column_id(id)
        self.devicesView.append_column(column)
    
    def add_toggle_column(self, title, id):
        toggle = gtk.CellRendererToggle()
        toggle.set_property('activatable', True)
        toggle.connect('toggled', self.mount_check)
        column = gtk.TreeViewColumn(title, toggle, active=id)
        self.devicesView.append_column(column)

    def mount_check(self, cell, path):
        for device in self.devices:
            if device[0] == self.devicesList[path][0]:
                if self.devicesList[path][2]:
                    device[1].umount()
                else:
                    device[1].mount()
        self.refresh()
    
    def row_click(self, widget, column = None, pointer = None):
        selection = self.devicesView.get_selection()
        model, selection_iter = selection.get_selected()
        if not selection_iter:
            return
        
        # Get label
        label = model.get_value(selection_iter,0)
        # Get device
        device = self.devices.get_device(label)
        device.get_info()
        
        self.devices.properties = PropertiesWindow(device, self.devices, True)
        self.devices.properties.run()
        
    def refresh(self):
        self.devices.list_all_devices()
        self.devicesList.clear()
        for name, device in self.devices:
            self.devicesList.append([name, device.block_device, device.is_mounted])
            
    def refresh_properties(self):
        self.devicesList.clear()
        for name, device in self.devices:
            device.get_info()
            name = device.label
            self.devicesList.append([name, device.block_device, device.is_mounted])
    
    def destroy(self, widget, data=None):
        self.devices.main = None
        self.window.hide()
        self.window.destroy()
        gtk.main_quit()
    
    def run(self):
        gtk.main()
    

if __name__ == '__main__':
    #argv.append('/dev/sd1')
    if len(argv) is 1:
        DBusGMainLoop(set_as_default=True)
        usbdevices = USBDevices()
        usbdevices.run()
        
    elif len(argv) is 2:
        DBusGMainLoop(set_as_default=True)
        usbdevices = USBDevices()
        usbdevices.run(argv[1])
        
    else:
        print "Usage:"
        print "   usbmanager [mount_point]"
        print "ex:"
        print "   $ usbmanager"
        print "or "
        print "   $ usbmanager /media/my_flash_device"
        exit(0)
    exit(1)
